# -- coding: utf-8 --
import threading
import msvcrt
import ctypes
import time
import os,cv2
import numpy as np
from ctypes import *
from Mv3dRgbdImport.Mv3dRgbdDefine import *
from Mv3dRgbdImport.Mv3dRgbdApi import *
from Mv3dRgbdImport.Mv3dRgbdDefine import DeviceType_Ethernet, DeviceType_USB, DeviceType_Ethernet_Vir, DeviceType_USB_Vir, MV3D_RGBD_FLOAT_EXPOSURETIME, \
    ParamType_Float, ParamType_Int, ParamType_Enum, CoordinateType_Depth, MV3D_RGBD_FLOAT_Z_UNIT

def get_depth_at_point(depth_image, x, y):
    """获取指定坐标的深度值"""
    if depth_image is None:
        print("深度图像未初始化")
        return None
    
    if x < 0 or x >= depth_image.shape[1] or y < 0 or y >= depth_image.shape[0]:
        print("坐标超出图像范围")
        return None
    
    depth_value = depth_image[y, x]
    print(f"坐标({x}, {y})处的深度值为: {depth_value} mm")
    return depth_value

def yuv2rgb(np_color,height,width):
   # 方法1：直接转换（推荐）
    try:
        yuyv_image = np_color.reshape((height, width * 2))
        rgb_image = cv2.cvtColor(yuyv_image, cv2.COLOR_YUV2BGR_YUY2)
    except Exception as e:
        # print(f"直接转换失败: {e}")
        
        # 方法2：精确手动转换
        # YUYV格式布局：Y0 U0 Y1 V0 Y2 U2 Y3 V2...
        # 总数据量：width*height*2 bytes
        
        # 创建完整数组
        yuyv_array = np_color.reshape(-1, 4)  # 每组4字节: Y0 U0 Y1 V0
        
        # 提取Y分量（所有像素）
        y = np.zeros(height * width, dtype=np.uint8)
        y[0::2] = yuyv_array[:, 0]  # 偶数像素Y
        y[1::2] = yuyv_array[:, 2]  # 奇数像素Y
        
        # 提取U和V分量（每两个像素共享一组UV）
        u = yuyv_array[:, 1]
        v = yuyv_array[:, 3]
        
        # 调整UV分量尺寸以匹配Y分量
        u = np.repeat(u, 2)  # 将U值重复给两个像素
        v = np.repeat(v, 2)  # 将V值重复给两个像素
        
        # 重组为YUV格式（需要转置为OpenCV期望的格式）
        y = y.reshape(height, width)
        u = u.reshape(height, width)
        v = v.reshape(height, width)
        
        # 合并为3通道图像
        yuv_image = np.dstack((y, u, v)).astype(np.uint8)
        
        # 转换为BGR
        rgb_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR)
    
    return rgb_image

def parse_frame(stFrameData):
    if stFrameData.nImageCount != 0:
        color_img = None
        depth_img = None
        
        for i in range(0, stFrameData.nImageCount):
            # print(f"MV3D_RGBD_FetchFrame[{i}]: nFrameNum[{stFrameData.stImageData[i].nFrameNum}], "
            #       f"nDataLen[{stFrameData.stImageData[i].nDataLen}], "
            #       f"nWidth[{stFrameData.stImageData[i].nWidth}], "
            #       f"nHeight[{stFrameData.stImageData[i].nHeight}]")
            
            # print('矫正：', stFrameData.stImageData[i].bIsRectified)
            # print('图像格式：', stFrameData.stImageData[i].enImageType)
            
            # 处理深度图像
            if stFrameData.stImageData[i].enImageType == 17825976:
                depth_data = stFrameData.stImageData[i].pData
                height = stFrameData.stImageData[i].nHeight
                width = stFrameData.stImageData[i].nWidth
                
                # 将深度数据转换为numpy数组
                np_depth = np.frombuffer((c_ubyte * stFrameData.stImageData[i].nDataLen).from_address(
                    ctypes.addressof(depth_data.contents)), dtype=np.uint16)
                np_depth = np_depth.reshape((height, width))
                
                # 保存深度图像
                depth_img = np_depth.copy()
                global_depth_image = depth_img
                
                # # 显示深度图像
                # depth_vis = cv2.normalize(np_depth, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
                # cv2.imshow('Depth Image', depth_vis)
            
            # 处理彩色图像 (YUV422格式)
            elif stFrameData.stImageData[i].enImageType == 34603058:
                color_data = stFrameData.stImageData[i].pData
                height = stFrameData.stImageData[i].nHeight
                width = stFrameData.stImageData[i].nWidth
                
                # 将YUV422数据转换为numpy数组
                np_color = np.frombuffer((c_ubyte * stFrameData.stImageData[i].nDataLen).from_address(
                    ctypes.addressof(color_data.contents)), dtype=np.uint8)
                
                # # YUV422转RGB
                # yuv_image = np_color.reshape((height, width * 2))
                # # rgb_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR_YUYV)
                # yuv420_image = cv2.resize(yuv_image, (width, height*2))
                # rgb_image = cv2.cvtColor(yuv420_image, cv2.COLOR_YUV2BGR_I420)
                rgb_image=yuv2rgb(np_color,height,width)
                
                # 保存彩色图像
                color_img = rgb_image.copy()
                global_color_image = color_img
    return color_img,depth_img

def save_images(color_img, depth_img, prefix="frame"):
    """保存彩色和深度图像"""
    if color_img is not None:
        color_path = f"{prefix}_color.png"
        cv2.imwrite(color_path, color_img)
        print(f"已保存彩色图像到: {color_path}")
    
    if depth_img is not None:
        # 归一化深度图像以便保存
        depth_normalized = cv2.normalize(depth_img, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
        depth_path = f"{prefix}_depth.png"
        cv2.imwrite(depth_path, depth_normalized)
        print(f"已保存深度图像到: {depth_path}")



g_bExit = False
class HKCamera:
    def __init__(self):
        self.camera=None
        self.init_hkcamera()
    def init_hkcamera(self):
        #初始化摄像头
        nDeviceNum=ctypes.c_uint(0)
        nDeviceNum_p=byref(nDeviceNum)
        # ch:获取设备数量 | en:Get device number
        ret=Mv3dRgbd.MV3D_RGBD_GetDeviceNumber(DeviceType_Ethernet | DeviceType_USB | DeviceType_Ethernet_Vir | DeviceType_USB_Vir, nDeviceNum_p) 
        if  ret!=0:
            print("MV3D_RGBD_GetDeviceNumber fail! ret[0x%x]" % ret)
            os.system('pause')
            sys.exit()
        if  nDeviceNum==0:
            print("find no device!")
            os.system('pause')
            sys.exit()
        print("Find devices numbers:", nDeviceNum.value)
        
        stDeviceList = MV3D_RGBD_DEVICE_INFO_LIST()
        net = Mv3dRgbd.MV3D_RGBD_GetDeviceList(DeviceType_Ethernet | DeviceType_USB | DeviceType_Ethernet_Vir | DeviceType_USB_Vir, pointer(stDeviceList.DeviceInfo[0]), 20, nDeviceNum_p)
        
        for i in range(0, nDeviceNum.value):
            print("\ndevice: [%d]" % i)
            strModeName = ""
            for per in stDeviceList.DeviceInfo[i].chModelName:
                strModeName = strModeName + chr(per)
            print("device model name: %s" % strModeName)

            strSerialNumber = ""
            for per in stDeviceList.DeviceInfo[i].chSerialNumber:
                strSerialNumber = strSerialNumber + chr(per)
            print("device SerialNumber: %s" % strSerialNumber)

        # ch:创建相机示例 | en:Create a camera instance
        self.camera=Mv3dRgbd()
        # nConnectionNum = input("please input the number of the device to connect:")
        nConnectionNum = 0
        if int(nConnectionNum) >= nDeviceNum.value:
            print("intput error!")
            os.system('pause')
            sys.exit()

        # ch:打开设备 | en:Open device  
        ret = self.camera.MV3D_RGBD_OpenDevice(pointer(stDeviceList.DeviceInfo[int(nConnectionNum)]))
        if ret != 0:
            print ("MV3D_RGBD_OpenDevice fail! ret[0x%x]" % ret)
            os.system('pause')
            sys.exit()
    
    def destroy_camera(self):
        # ch:停止取流 | en:Stop grabbing
        if self.camera is not None:
            ret=self.camera.MV3D_RGBD_Stop()
            if ret != 0:
                print ("stop fail! ret[0x%x]" % ret)
                os.system('pause')
                sys.exit()

            # ch:销毁句柄 | en:Destroy the device handle 
            ret=self.camera.MV3D_RGBD_CloseDevice()
            if ret != 0:
                print ("CloseDevice fail! ret[0x%x]" % ret)
                os.system('pause')
                sys.exit()

    def work_thread(self,pdata=0,nDataSize=0):
        # ch:开始取流 | en:Start grabbing
        ret=self.camera.MV3D_RGBD_Start()
        if ret != 0:
            print ("start fail! ret[0x%x]" % ret)
            self.camera.MV3D_RGBD_CloseDevice()
            os.system('pause')
            sys.exit()
        while True:
            stFrameData=MV3D_RGBD_FRAME_DATA()
            ret=self.camera.MV3D_RGBD_FetchFrame(pointer(stFrameData), 5000)
            if ret==0:
                color_img,depth_img = parse_frame(stFrameData)
                save_images(color_img,depth_img)
            else:
                print("no data[0x%x]" % ret)
            if g_bExit == True:
                break

if __name__ == "__main__":


    stvison = HKCamera()

    # ch:获取图像线程 | en:Get image thread
    try:
        hthreadhandle=threading.Thread(target=stvison.work_thread,args=(None,None))
        hthreadhandle.start()
    except:
        print("error: unable to start thread")

    os.system('pause')
    g_bExit = True



